/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.bff.gaia.optimizer.traversals;

import com.bff.gaia.api.common.functions.ProjectFlatJoinFunction;
import com.bff.gaia.optimizer.DataStatistics;
import com.bff.gaia.optimizer.dag.*;
import com.bff.gaia.util.Visitor;

import java.util.*;

/**
 * This traversal of the optimizer DAG assigns IDs to each node (in a pre-order fashion),
 * and calls each node to compute its estimates. The latter happens in the postVisit function,
 * where it is guaranteed that all predecessors have computed their estimates.
 */
public class IdAndEstimatesVisitor implements Visitor<OptimizerNode> {

	private final DataStatistics statistics;

	private int id = 1;

	public ArrayList<JoinTreeRecord> Record = new ArrayList<>();

	public IdAndEstimatesVisitor(DataStatistics statistics) {
		this.statistics = statistics;
	}

	@Override
	public boolean preVisit(OptimizerNode visitable) {
		return visitable.getId() == -1;
	}

	@Override
	public void postVisit(OptimizerNode visitable) {
		// the node ids
		visitable.initId(this.id++);

		if (visitable instanceof JoinTreeNode) {
			addRecord((JoinTreeNode) visitable);
		}

		// connections need to figure out their maximum path depths
		for (DagConnection conn : visitable.getIncomingConnections()) {
			conn.initMaxDepth();
		}
		for (DagConnection conn : visitable.getBroadcastConnections()) {
			conn.initMaxDepth();
		}

		// the estimates
		visitable.computeOutputEstimates(this.statistics);

		// if required, recurse into the step function
		if (visitable instanceof IterationNode) {
			((IterationNode) visitable).acceptForStepFunction(this);
		}
	}

	/**
	 * 记录待优化节点
	 */
	public class JoinTreeRecord {

		private int joinTreeNodeId = -1;
		private int joinTreeSourceNodeId = -1;
		private final Map<Integer, OptimizerNode> joinTreeNodeMap = new HashMap<>();
		private final Map<Integer, OptimizerNode> joinTreeSourceNodeMap = new HashMap<>();

		private ArrayList<ColumnSourceNodeSet> keySetList = new ArrayList<>();
		private ArrayList<ColumnSourceNodeSet> resultSetList = new ArrayList<>();

		public class ColumnSourceNodeSet {

			private int id;
			private boolean finished = false;
			private Set<ColumnSourceNode> columnSourceNodeSet;

			ColumnSourceNodeSet(Set<ColumnSourceNode> ColumnSourceNodes) {
				this.columnSourceNodeSet = ColumnSourceNodes;
			}

			public int getId() {
				return id;
			}

			private void setId(int id) {
				this.id = id;
			}

			public boolean checkFinished() {
				finished = true;
				for (ColumnSourceNode columnSourceNode : columnSourceNodeSet) {
					finished = finished && columnSourceNode.isFinished();
				}
				return finished;
			}

			public boolean isFinished() {
				return finished;
			}

			public void setFinished(boolean finished) {
				this.finished = finished;
			}

			public Set<ColumnSourceNode> getColumnSourceNodeSet() {
				return columnSourceNodeSet;
			}

			boolean containNode(int id, int pos) {
				return columnSourceNodeSet.contains(new ColumnSourceNode(id, pos));
			}

			boolean containNode(int id) {
				for (ColumnSourceNode node : columnSourceNodeSet) {
					if (node.getNodeId() == id) {
						return true;
					}
				}
				return false;
			}

			boolean intersectSet(HashSet<Integer> set) {
				for (int id : set) {
					if (this.containNode(id)) {
						return true;
					}
				}
				return false;
			}
		}

		public class ColumnSourceNode {

			private int nodeId;
			private int superId;
			private boolean finished = false;
			private ArrayList<Integer> columnPos = new ArrayList<>();

			ColumnSourceNode(int id, int keyPos) {
				this.nodeId = id;
				this.columnPos.add(keyPos);
			}

			public int getNodeId() {
				return nodeId;
			}

			public int getSuperId() {
				return superId;
			}

			private void setSuperId(int id) {
				this.superId = id;
			}

			public boolean isFinished() {
				return finished;
			}

			public void setFinished(boolean finished) {
				this.finished = finished;
			}

			public ArrayList<Integer> getColumnPos() {
				return columnPos;
			}

			private ColumnSourceNode isBelongTo(ColumnSourceNodeSet set) {
				for (ColumnSourceNode node : set.getColumnSourceNodeSet()) {
					if (node.equals(new ColumnSourceNode(this.nodeId, this.columnPos.get(0)))) {
						return node;
					}
				}
				return null;
			}

			@Override
			public boolean equals(Object o) {
				if (this == o) return true;
				if (o == null || getClass() != o.getClass()) return false;

				ColumnSourceNode columnSourceNode = (JoinTreeRecord.ColumnSourceNode) o;

				return nodeId == columnSourceNode.nodeId && columnPos.equals(columnSourceNode.columnPos);
			}

			@Override
			public int hashCode() {
				int result = nodeId;
				result = 31 * result + columnPos.hashCode();
				return result;
			}
		}

		/**
		 * @param node      要溯源的节点
		 * @param columnPos 要溯源的列
		 * @return 包含该列的所有Source节点集合
		 */
		private Set<ColumnSourceNode> trace(OptimizerNode node, int columnPos) {

			Set<ColumnSourceNode> sourceNodeSet = new HashSet<>();

			if (node instanceof DataSourceNode) {
				sourceNodeSet.add(new ColumnSourceNode(node.getJoinTreeSourceNodeId(), columnPos));
			} else if (node instanceof JoinTreeNode) {
				List<DagConnection> joinInput = node.getIncomingConnections();
				int keyFrom[] = {((JoinTreeNode) node).getOperator().getKeyColumns(0)[0], ((JoinTreeNode) node).getOperator().getKeyColumns(1)[0]};
				int i;
				for (i = 0; i < 2; i++) {
					int forwardColumnPos = node.getSemanticProperties().getForwardingSourceField(i, columnPos);
					if (forwardColumnPos != -1) {
						sourceNodeSet.addAll(trace(joinInput.get(i).getSource(), forwardColumnPos));
						if (forwardColumnPos == keyFrom[i]) {
							sourceNodeSet.addAll(trace(joinInput.get(1 - i).getSource(), keyFrom[1 - i]));
							break;
						}
						break;
					}
				}
			}
			return sourceNodeSet;
		}

		/**
		 * 寻找所有显性和隐性的连接关系
		 */
		public void initKeySetList() {

			for (OptimizerNode node : joinTreeNodeMap.values()) {
				int keyFrom[] = {((JoinTreeNode) node).getOperator().getKeyColumns(0)[0], ((JoinTreeNode) node).getOperator().getKeyColumns(1)[0]};
				Set<ColumnSourceNode> keyNodeSet = new HashSet<>();
				keyNodeSet.addAll(trace(((JoinTreeNode) node).getFirstPredecessorNode(), keyFrom[0]));
				keyNodeSet.addAll(trace(((JoinTreeNode) node).getSecondPredecessorNode(), keyFrom[1]));
				boolean isNew = true;
				ColumnSourceNodeSet p = null;
				Iterator<ColumnSourceNodeSet> iterator = keySetList.iterator();
				while (iterator.hasNext()) {
					ColumnSourceNodeSet keySet = iterator.next();
					for (ColumnSourceNode keyNode : keyNodeSet) {
						if (keySet.columnSourceNodeSet.contains(keyNode)) {
							isNew = false;
							break;
						}
					}
					if (!isNew) {
						if (p == null) {
							keySet.columnSourceNodeSet.addAll(keyNodeSet);
							p = keySet;
						} else {
							p.columnSourceNodeSet.addAll(keyNodeSet);
							iterator.remove();
						}
					}
				}
				if (isNew) {
					keySetList.add(new ColumnSourceNodeSet(keyNodeSet));
				}
			}
			initId(keySetList);
		}

		public void initResultSetList() {
			JoinTreeNode tailJoinTreeNode = (JoinTreeNode) joinTreeNodeMap.get(joinTreeNodeMap.size() - 1);
			int resultLength = ((ProjectFlatJoinFunction) tailJoinTreeNode.getOperator().getUserCodeWrapper().getUserCodeObject()).getFields().length;
			for (int i = 0; i < resultLength; i++) {
				ColumnSourceNodeSet resultNodeSet = new ColumnSourceNodeSet(trace(tailJoinTreeNode, i));
				resultSetList.add(resultNodeSet);
			}
			initId(resultSetList);
		}

		private void initId(ArrayList<ColumnSourceNodeSet> setList) {
			for (int i = 0; i < setList.size(); i++) {
				ColumnSourceNodeSet columnSourceNodeSet = setList.get(i);
				columnSourceNodeSet.setId(i);
				for (ColumnSourceNode columnSourceNode : columnSourceNodeSet.getColumnSourceNodeSet()) {
					columnSourceNode.setSuperId(i);
				}
			}
		}

		public void setNodeFinished(ArrayList<ColumnSourceNodeSet> setList, int id) {
			for (ColumnSourceNodeSet columnSourceNodeSet : setList) {
				for (ColumnSourceNode columnSourceNode : columnSourceNodeSet.getColumnSourceNodeSet()) {
					if (columnSourceNode.getNodeId() == id) {
						columnSourceNode.setFinished(true);
					}
				}
			}
		}

		private void setNodeFinished(ArrayList<ColumnSourceNodeSet> setList, HashSet<Integer> set) {
			for (int id : set) {
				setNodeFinished(setList, id);
			}
		}

		private void resetNodeFinished(ArrayList<ColumnSourceNodeSet> setList) {
			for (ColumnSourceNodeSet columnSourceNodeSet : setList) {
				for (ColumnSourceNode columnSourceNode : columnSourceNodeSet.getColumnSourceNodeSet()) {
					columnSourceNode.setFinished(false);
				}
			}
		}

		public int width(HashSet<Integer> set) {
			int width = 0;
			setNodeFinished(keySetList, set);
			for (ColumnSourceNodeSet keySet : keySetList) {
				if (!keySet.checkFinished() && keySet.intersectSet(set)) {
					width++;
				}
			}
			for (ColumnSourceNodeSet resultSet : resultSetList) {
				label:
				if (resultSet.intersectSet(set)) {
					for (ColumnSourceNode resultNode : resultSet.getColumnSourceNodeSet()) {
						for (ColumnSourceNodeSet columnSourceNodeSet : keySetList) {
							ColumnSourceNode node = resultNode.isBelongTo(columnSourceNodeSet);
							if (node == null || keySetList.get(node.getSuperId()).isFinished()) {
								width++;
								break label;
							}
						}
					}
				}
			}
			resetNodeFinished(keySetList);
			return width;
		}

		public int getSuperSetId(ArrayList<ColumnSourceNodeSet> setList, int nodeId, int pos) {
			for (ColumnSourceNodeSet columnSourceNodeSet : setList) {
				if (columnSourceNodeSet.containNode(nodeId, pos)) {
					return columnSourceNodeSet.getId();
				}
			}
			return -1;
		}

		boolean isConnectable(int nodeId1, int nodeId2) {
			for (ColumnSourceNodeSet keySet : keySetList) {
				if (keySet.containNode(nodeId1) && keySet.containNode(nodeId2)) {
					return true;
				}
			}
			return false;
		}

		public boolean isConnectable(HashSet<Integer> set1, HashSet<Integer> set2) {
			for (int nodeId1 : set1) {
				for (int nodeId2 : set2) {
					if (isConnectable(nodeId1, nodeId2)) {
						return true;
					}
				}
			}
			return false;
		}

		public ArrayList<ColumnSourceNodeSet> getKeySetList() {
			return keySetList;
		}

		public ArrayList<ColumnSourceNodeSet> getResultSetList() {
			return resultSetList;
		}

		public Map<Integer, OptimizerNode> getJoinTreeNodeMap() {
			return joinTreeNodeMap;
		}

		public Map<Integer, OptimizerNode> getJoinTreeSourceNodeMap() {
			return joinTreeSourceNodeMap;
		}
	}

	private void addRecord(JoinTreeNode visitable) {

		List<DagConnection> incomingConnections = visitable.getIncomingConnections();

		boolean isNew = true;

		for (DagConnection joinInput : incomingConnections) {
			if (joinInput.getSource() instanceof JoinTreeNode) {
				isNew = false;
			}
		}

		if (isNew) {
			Record.add(new JoinTreeRecord());
		}

		JoinTreeRecord tailRecord = Record.get(Record.size() - 1);
		tailRecord.joinTreeNodeMap.put(++tailRecord.joinTreeNodeId, visitable);

		for (DagConnection joinInput : incomingConnections) {
			if (!(joinInput.getSource() instanceof JoinTreeNode)) {
				joinInput.getSource().initJoinTreeSourceNodeId(++tailRecord.joinTreeSourceNodeId);
				tailRecord.joinTreeSourceNodeMap.put(tailRecord.joinTreeSourceNodeId, joinInput.getSource());
			}
		}
	}
}